package com.eudon.ai.agent.etl.transformer;

import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;

/**
 * @author eudon
 * @description 关键词增强器，在meta信息中新增关键词
 * @date 2025/8/15
 * @email syd19970616@gmail.com
 * @address <a href="https://gitee.com/eudon">Gitee</a>
 */
@Slf4j
@Component
public class MyKeywordEnricher {

    public static final String CONTEXT_STR_PLACEHOLDER = "context";

    public static final String KEYWORDS_TEMPLATE = """
            {context}. Give %s unique keywords for this
            document. Format as comma separated. Keywords: """;

    private static final String EXCERPT_KEYWORDS_METADATA_KEY = "excerpt_keywords";
    @Resource
    private ChatModel dashscopeChatModel;

    public List<Document> addKeyword(List<Document> documents, int keywordCount, String keywordName) {
        //调用阿里云百炼模型根据每一条document提取5个关键词，添加到document元数据中
//        KeywordMetadataEnricher keywordMetadataEnricher = new KeywordMetadataEnricher(dashScopeChatModel, 5);
        return this.apply(documents, keywordCount, keywordName);
    }

    public List<Document> apply(List<Document> documents, int keywordCount, String keywordName) {
        for (Document document : documents) {

            var template = new PromptTemplate(String.format(KEYWORDS_TEMPLATE, keywordCount));
            Prompt prompt = template.create(Map.of(CONTEXT_STR_PLACEHOLDER, document.getText()));
            String keywords = dashscopeChatModel.call(prompt).getResult().getOutput().getText();
            document.getMetadata().put(StringUtils.isBlank(keywordName) ? EXCERPT_KEYWORDS_METADATA_KEY : keywordName, keywords);
        }
        return documents;
    }
}
